package com.jvmmonitor.util;

import java.lang.instrument.Instrumentation;
import java.lang.reflect.Array;
import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.util.IdentityHashMap;
import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 对象大小计算工具类
 */
public class ObjectSizeCalculator {
    private static final Logger logger = LoggerFactory.getLogger(ObjectSizeCalculator.class);
    
    // Instrumentation实例，用于精确计算对象大小
    private static Instrumentation instrumentation;
    
    // 对象头大小（64位JVM，未开启压缩指针）
    private static final int OBJECT_HEADER_SIZE = 16;
    
    // 数组头大小（64位JVM，未开启压缩指针）
    private static final int ARRAY_HEADER_SIZE = 24;
    
    // 引用大小（64位JVM，未开启压缩指针）
    private static final int REFERENCE_SIZE = 8;
    
    // 基本类型大小
    private static final Map<Class<?>, Integer> PRIMITIVE_SIZES = new IdentityHashMap<>();
    
    static {
        // 初始化基本类型大小
        PRIMITIVE_SIZES.put(boolean.class, 1);
        PRIMITIVE_SIZES.put(byte.class, 1);
        PRIMITIVE_SIZES.put(char.class, 2);
        PRIMITIVE_SIZES.put(short.class, 2);
        PRIMITIVE_SIZES.put(int.class, 4);
        PRIMITIVE_SIZES.put(float.class, 4);
        PRIMITIVE_SIZES.put(long.class, 8);
        PRIMITIVE_SIZES.put(double.class, 8);
    }
    
    /**
     * 设置Instrumentation实例
     */
    public static void setInstrumentation(Instrumentation inst) {
        instrumentation = inst;
        logger.info("Instrumentation set for object size calculation");
    }
    
    /**
     * 计算对象大小
     * @param obj 要计算大小的对象
     * @return 对象大小（字节）
     */
    public static long calculateObjectSize(Object obj) {
        if (obj == null) {
            return 0;
        }
        
        // 如果有Instrumentation实例，使用它来精确计算
        if (instrumentation != null) {
            try {
                return instrumentation.getObjectSize(obj);
            } catch (Exception e) {
                logger.warn("Failed to calculate object size using Instrumentation: {}, falling back to estimation", e.getMessage());
            }
        }
        
        // 否则使用估算方法
        return estimateObjectSize(obj, new IdentityHashMap<>());
    }
    
    /**
     * 估算对象大小（递归）
     * @param obj 要计算大小的对象
     * @param visited 已访问对象集合，防止循环引用
     * @return 对象大小估算值（字节）
     */
    private static long estimateObjectSize(Object obj, Map<Object, Object> visited) {
        if (obj == null || visited.containsKey(obj)) {
            return 0;
        }
        
        visited.put(obj, null);
        Class<?> clazz = obj.getClass();
        
        // 基本类型包装类
        if (clazz.isPrimitive()) {
            return PRIMITIVE_SIZES.getOrDefault(clazz, 0);
        }
        
        // 字符串特殊处理
        if (clazz == String.class) {
            String str = (String) obj;
            return OBJECT_HEADER_SIZE + 4 + 4 + (str.length() * 2);
        }
        
        // 数组处理
        if (clazz.isArray()) {
            long size = ARRAY_HEADER_SIZE;
            int length = Array.getLength(obj);
            
            if (length > 0) {
                Class<?> componentType = clazz.getComponentType();
                if (componentType.isPrimitive()) {
                    size += length * PRIMITIVE_SIZES.getOrDefault(componentType, 0);
                } else {
                    size += length * REFERENCE_SIZE;
                    // 递归计算数组元素大小
                    for (int i = 0; i < length; i++) {
                        Object element = Array.get(obj, i);
                        size += estimateObjectSize(element, visited);
                    }
                }
            }
            
            return alignTo8Bytes(size);
        }
        
        // 普通对象处理
        long size = OBJECT_HEADER_SIZE;
        
        // 遍历所有字段（包括父类字段）
        for (Class<?> currentClass = clazz; currentClass != null; currentClass = currentClass.getSuperclass()) {
            Field[] fields = currentClass.getDeclaredFields();
            for (Field field : fields) {
                // 跳过静态字段
                if (Modifier.isStatic(field.getModifiers())) {
                    continue;
                }
                
                Class<?> fieldType = field.getType();
                if (fieldType.isPrimitive()) {
                    size += PRIMITIVE_SIZES.getOrDefault(fieldType, 0);
                } else {
                    size += REFERENCE_SIZE;
                    // 递归计算引用对象大小
                    try {
                        field.setAccessible(true);
                        Object fieldValue = field.get(obj);
                        size += estimateObjectSize(fieldValue, visited);
                    } catch (IllegalAccessException e) {
                        logger.debug("Cannot access field {}: {}", field.getName(), e.getMessage());
                    }
                }
            }
        }
        
        return alignTo8Bytes(size);
    }
    
    /**
     * 对齐到8字节边界
     * @param size 原始大小
     * @return 对齐后的大小
     */
    private static long alignTo8Bytes(long size) {
        return (size + 7) & ~7;
    }
}